# -*- coding: utf-8 -*-
"""
Created on Mon Apr  1 22:01:03 2024
given constraints and cost function in terms of phi, B, coil(etc), find coil current Ia

@author: xiang
"""
import numpy as np
import math
import matplotlib.pyplot as plt

from scipy import special

import multiprocessing as mp
from calculate_green_function import green_function as GF


class fix_boundary_gs():
    
    
    def __init__(self):
        # self.machine_name = settings.MACHINE_NAME
        # self.data_path = settings.DATA_PATH
        # self.machine_file = settings.MACHINE_FILE
        self.machine_file = 'C:\\Users\\xiang\\Documents\\XINAO\\AI\\SEPY\\data\\EXL50U\\machine.xlsx'
        # self.machine_file = 'C:\\sepy\\data\\EXL50U\\machine.xlsx'
        self.cmu = np.pi*2.0e-7;
        # self.num_cores = int(mp.cpu_count())
        eq = np.load('C:\\Users\\xiang\\Documents\\XINAO\\AI\\SEPY\\data\\EXL50U\\green_function.npy', 
                      allow_pickle=True).item()
        
        # eq = np.load('C:\\sepy\\data\\EXL50U\\green_function.npy', 
        #                 allow_pickle=True).item()
        
        # fgs dictionary is used for free gs solve
        self.fgs = {}
        
        # default region for finding x point 
        self.fgs['index_xpoint_up_x']   = 16
        self.fgs['index_xpoint_up_y']   = 20
        self.fgs['index_xpoint_down_x'] = 113
        self.fgs['index_xpoint_down_y'] = 20
        self.fgs['index_xpoint_height'] = 20
        self.fgs['index_xpoint_width']  = 20
        
        
        self.eq = eq
                
        self.num_cores = int(mp.cpu_count())
        

    
    
    def init_plasma(self, isplot=0):
        """
        initial plasma for intergration 
        
        """
        
        # green funciton 
        mi_plasma_pf = self.eq['green']['mi_plasma_pf']
        mi_limiter_pf = self.eq['green']['mi_limiter_pf']
        
        # plasma and limiter region
        x_plasma = self.eq['green']['grid']['plasma']['x']
        y_plasma = self.eq['green']['grid']['plasma']['y']
        num_plasma = len(x_plasma)
        gapx = (x_plasma[1]-x_plasma[0])/2
        
        
        
        # flux of plasma
        # flux inducted by plasma on plasma
        I_init, x_init, y_init = self.get_init_plasma_para()
        num_init = len(x_init)
        # flux_plasma_plasma = self.calculate_mi_plasma_init(x_plasma, y_plasma, 
        #                                                    x_init, y_init, I_init)
        num_init = len(I_init)
        
        flux_plasma_init = np.zeros(shape=(num_plasma, num_init))
        # pool = mp.Pool(self.num_cores)        
        # flux_plasma_init = [pool.apply(self.calculate_mi_plasma_init, 
        #                                args=(x_plasma, y_plasma, x_init[i], y_init[i], I_init[i], gapx)) 
        #                     for i in range(num_init)]        
        for i in range(num_init):
            flux_plasma_init[:,i] = self.calculate_mi_plasma_init(x_plasma, y_plasma, 
                                                                  x_init[i], y_init[i], 
                                                                  I_init[i], gapx)
        flux_plasma_init = np.sum(flux_plasma_init, axis=1)
        
        # flux inducted by PF current on plasma
        flux_plasma_pfs = pfs.dot(mi_plasma_pf)
        flux_plasma_pfs = flux_plasma_pfs.reshape(x_plasma.shape)
        
        # flux on plasma
        flux_plasma = flux_plasma_init + flux_plasma_pfs
        
        
        
        # flux of limiter
        x_limiter = self.eq['machine']['limiter']['x_limiter']
        y_limiter = self.eq['machine']['limiter']['y_limiter']
        num_limiter = len(x_limiter)
        
        # flux inducted on limiter by pf
        flux_limiter_pfs = pfs.dot(mi_limiter_pf)
        
        # flux inducted on limiter by init plasma        
        flux_limiter_init = np.zeros(shape=(num_limiter, num_init));
        for i in range(num_init):
            flux_limiter_init[:,i] = self.calculate_mi_limiter_init(x_limiter, y_limiter, 
                                                                  x_init[i], y_init[i], 
                                                                  I_init[i], gapx)
        flux_limiter_init = np.sum(flux_limiter_init, axis=1)
        
        # flux inducted on limiter
        flux_limiter = flux_limiter_pfs + flux_limiter_init
        
        
        
        # 3: identify plasma of init
        # identify plasma with plasma and limiter flux
        flux_plasma, flux_plasma_norm, index = self.identity_plasma_init(flux_plasma, flux_limiter, init_factor=10)
        
        
        
        # calculate plasma with phi
        j_plasma = self.iterate_plasma_with_phi(flux_plasma, flux_plasma_norm, index)        
        
        
        #
        if isplot:
            plt.figure()
            plt.plot(x_init, y_init, '.')
        
        return j_plasma, flux_plasma

    def generate_boundary_condition(self, kappa=2, aspect_ratio=1.8, R=1.05, delta=0.5):
        
        


    def iterate_phi_with_currents(self, j_plasma, pfs):
        
            flux_plasma, flux_plasma_norm, index = [], [], []
            
        return flux_plasma, flux_plasma_norm, index



    def iterate_pfs_with_phi(self, flux_plasma, flux_plasma_norm, index)

        j_plasma, pfs = [], []

        return j_plasma, pfs



    def fix_boundary_iterate(self, plasma, boundary_condition, n_iterate):
        
        j_plasma, flux_plasma, pfs = self.init_plasma()       
        
        for i in range(n_iterate):
            print(i)
            
            # flux induced by pf current and plasma current 
            flux_plasma, flux_plasma_norm, index = self.iterate_phi_with_currents(j_plasma, pfs)
            
            # find difference between given phi and boundary condition, get changed pf current
            j_plasma, pfs = self.iterate_pfs_with_phi(flux_plasma, flux_plasma_norm, index)
        
            self.plot_flux(flux_plasma, islcfs=1)
            
        return flux_plasma, j_plasma






if __name__ == '__main__':

    
    FBT = fix_boundary_gs()
    

    boundary_condition = FBT.generate_boundary_condition(kappa=2, aspect_ratio=1.8, R=1.05, delta=0.5)



    















